/*
 * PLyCursor类
 *
 * src/pl/plpython/plpy_cursorobject.c
 */

#include "postgres.h"

#include <limits.h>

#include "access/xact.h"
#include "catalog/pg_type.h"
#include "mb/pg_wchar.h"
#include "plpy_cursorobject.h"
#include "plpy_elog.h"
#include "plpy_main.h"
#include "plpy_planobject.h"
#include "plpy_procedure.h"
#include "plpy_resultobject.h"
#include "plpy_spi.h"
#include "plpython.h"
#include "utils/memutils.h"

static PyObject *fc_PLy_cursor_query(const char *fc_query);
static void fc_PLy_cursor_dealloc(PyObject *fc_arg);
static PyObject *fc_PLy_cursor_iternext(PyObject *fc_self);
static PyObject *fc_PLy_cursor_fetch(PyObject *fc_self, PyObject *fc_args);
static PyObject *fc_PLy_cursor_close(PyObject *fc_self, PyObject *fc_unused);

static char PLy_cursor_doc[] = "Wrapper around a PostgreSQL cursor";

static PyMethodDef PLy_cursor_methods[] = {
	{"fetch", fc_PLy_cursor_fetch, METH_VARARGS, NULL},
	{"close", fc_PLy_cursor_close, METH_NOARGS, NULL},
	{NULL, NULL, 0, NULL}
};

static PyTypeObject PLy_CursorType = {
	PyVarObject_HEAD_INIT(NULL, 0)
	.tp_name = "PLyCursor",
	.tp_basicsize = sizeof(PLyCursorObject),
	.tp_dealloc = fc_PLy_cursor_dealloc,
	.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
	.tp_doc = PLy_cursor_doc,
	.tp_iter = PyObject_SelfIter,
	.tp_iternext = fc_PLy_cursor_iternext,
	.tp_methods = PLy_cursor_methods,
};

void PLy_cursor_init_type(void)
{
	if (PyType_Ready(&PLy_CursorType) < 0)
		elog(ERROR, "could not initialize PLy_CursorType");
}

PyObject * PLy_cursor(PyObject *fc_self, PyObject *fc_args)
{
	char	   *fc_query;
	PyObject   *fc_plan;
	PyObject   *fc_planargs = NULL;

	if (PyArg_ParseTuple(fc_args, "s", &fc_query))
		return fc_PLy_cursor_query(fc_query);

	PyErr_Clear();

	if (PyArg_ParseTuple(fc_args, "O|O", &fc_plan, &fc_planargs))
		return PLy_cursor_plan(fc_plan, fc_planargs);

	PLy_exception_set(PLy_exc_error, "plpy.cursor expected a query or a plan");
	return NULL;
}


static PyObject * fc_PLy_cursor_query(const char *fc_query)
{
	PLyCursorObject *fc_cursor;
	PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();
	volatile MemoryContext fc_oldcontext;
	volatile ResourceOwner fc_oldowner;

	if ((fc_cursor = PyObject_New(PLyCursorObject, &PLy_CursorType)) == NULL)
		return NULL;
	fc_cursor->portalname = NULL;
	fc_cursor->closed = false;
	fc_cursor->mcxt = AllocSetContextCreate(TopMemoryContext,
										 "PL/Python cursor context",
										 ALLOCSET_DEFAULT_SIZES);

	/* 初始化以转换结果元组为Python */
	PLy_input_setup_func(&fc_cursor->result, fc_cursor->mcxt,
						 RECORDOID, -1,
						 fc_exec_ctx->curr_proc);

	fc_oldcontext = CurrentMemoryContext;
	fc_oldowner = CurrentResourceOwner;

	PLy_spi_subtransaction_begin(fc_oldcontext, fc_oldowner);

	PG_TRY();
	{
		SPIPlanPtr	fc_plan;
		Portal		fc_portal;

		pg_verifymbstr(fc_query, strlen(fc_query), false);

		fc_plan = SPI_prepare(fc_query, 0, NULL);
		if (fc_plan == NULL)
			elog(ERROR, "SPI_prepare failed: %s",
				 SPI_result_code_string(SPI_result));

		fc_portal = SPI_cursor_open(NULL, fc_plan, NULL, NULL,
								 fc_exec_ctx->curr_proc->fn_readonly);
		SPI_freeplan(fc_plan);

		if (fc_portal == NULL)
			elog(ERROR, "SPI_cursor_open() failed: %s",
				 SPI_result_code_string(SPI_result));

		fc_cursor->portalname = MemoryContextStrdup(fc_cursor->mcxt, fc_portal->name);

		PinPortal(fc_portal);

		PLy_spi_subtransaction_commit(fc_oldcontext, fc_oldowner);
	}
	PG_CATCH();
	{
		PLy_spi_subtransaction_abort(fc_oldcontext, fc_oldowner);
		return NULL;
	}
	PG_END_TRY();

	Assert(fc_cursor->portalname != NULL);
	return (PyObject *) fc_cursor;
}

PyObject * PLy_cursor_plan(PyObject *fc_ob, PyObject *fc_args)
{
	PLyCursorObject *fc_cursor;
	volatile int fc_nargs;
	PLyPlanObject *fc_plan;
	PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();
	volatile MemoryContext fc_oldcontext;
	volatile ResourceOwner fc_oldowner;

	if (fc_args)
	{
		if (!PySequence_Check(fc_args) || PyUnicode_Check(fc_args))
		{
			PLy_exception_set(PyExc_TypeError, "plpy.cursor takes a sequence as its second argument");
			return NULL;
		}
		fc_nargs = PySequence_Length(fc_args);
	}
	else
		fc_nargs = 0;

	fc_plan = (PLyPlanObject *) fc_ob;

	if (fc_nargs != fc_plan->nargs)
	{
		char	   *fc_sv;
		PyObject   *fc_so = PyObject_Str(fc_args);

		if (!fc_so)
			PLy_elog(ERROR, "could not execute plan");
		fc_sv = PLyUnicode_AsString(fc_so);
		PLy_exception_set_plural(PyExc_TypeError,
								 "Expected sequence of %d argument, got %d: %s",
								 "Expected sequence of %d arguments, got %d: %s",
								 fc_plan->nargs,
								 fc_plan->nargs, fc_nargs, fc_sv);
		Py_DECREF(fc_so);

		return NULL;
	}

	if ((fc_cursor = PyObject_New(PLyCursorObject, &PLy_CursorType)) == NULL)
		return NULL;
	fc_cursor->portalname = NULL;
	fc_cursor->closed = false;
	fc_cursor->mcxt = AllocSetContextCreate(TopMemoryContext,
										 "PL/Python cursor context",
										 ALLOCSET_DEFAULT_SIZES);

	/* 初始化以转换结果元组为Python */
	PLy_input_setup_func(&fc_cursor->result, fc_cursor->mcxt,
						 RECORDOID, -1,
						 fc_exec_ctx->curr_proc);

	fc_oldcontext = CurrentMemoryContext;
	fc_oldowner = CurrentResourceOwner;

	PLy_spi_subtransaction_begin(fc_oldcontext, fc_oldowner);

	PG_TRY();
	{
		Portal		fc_portal;
		MemoryContext fc_tmpcontext;
		Datum	   *volatile fc_values;
		char	   *volatile fc_nulls;
		volatile int fc_j;

		/*
		 * 转换后的参数和相关冗余信息将存在于这个上下文中，
		 * 该上下文是我们子事务的本地上下文。
		 */
		fc_tmpcontext = AllocSetContextCreate(CurTransactionContext,
										   "PL/Python temporary context",
										   ALLOCSET_SMALL_SIZES);
		MemoryContextSwitchTo(fc_tmpcontext);

		if (fc_nargs > 0)
		{
			fc_values = (Datum *) palloc(fc_nargs * sizeof(Datum));
			fc_nulls = (char *) palloc(fc_nargs * sizeof(char));
		}
		else
		{
			fc_values = NULL;
			fc_nulls = NULL;
		}

		for (fc_j = 0; fc_j < fc_nargs; fc_j++)
		{
			PLyObToDatum *fc_arg = &fc_plan->args[fc_j];
			PyObject   *fc_elem;

			fc_elem = PySequence_GetItem(fc_args, fc_j);
			PG_TRY();
			{
				bool		fc_isnull;

				fc_values[fc_j] = PLy_output_convert(fc_arg, fc_elem, &fc_isnull);
				fc_nulls[fc_j] = fc_isnull ? 'n' : ' ';
			}
			PG_FINALLY();
			{
				Py_DECREF(fc_elem);
			}
			PG_END_TRY();
		}

		MemoryContextSwitchTo(fc_oldcontext);

		fc_portal = SPI_cursor_open(NULL, fc_plan->plan, fc_values, fc_nulls,
								 fc_exec_ctx->curr_proc->fn_readonly);
		if (fc_portal == NULL)
			elog(ERROR, "SPI_cursor_open() failed: %s",
				 SPI_result_code_string(SPI_result));

		fc_cursor->portalname = MemoryContextStrdup(fc_cursor->mcxt, fc_portal->name);

		PinPortal(fc_portal);

		MemoryContextDelete(fc_tmpcontext);
		PLy_spi_subtransaction_commit(fc_oldcontext, fc_oldowner);
	}
	PG_CATCH();
	{
		Py_DECREF(fc_cursor);
		/* 子事务回滚将移除tmpcontext */
		PLy_spi_subtransaction_abort(fc_oldcontext, fc_oldowner);
		return NULL;
	}
	PG_END_TRY();

	Assert(fc_cursor->portalname != NULL);
	return (PyObject *) fc_cursor;
}

static void fc_PLy_cursor_dealloc(PyObject *fc_arg)
{
	PLyCursorObject *fc_cursor;
	Portal		fc_portal;

	fc_cursor = (PLyCursorObject *) fc_arg;

	if (!fc_cursor->closed)
	{
		fc_portal = GetPortalByName(fc_cursor->portalname);

		if (PortalIsValid(fc_portal))
		{
			UnpinPortal(fc_portal);
			SPI_cursor_close(fc_portal);
		}
		fc_cursor->closed = true;
	}
	if (fc_cursor->mcxt)
	{
		MemoryContextDelete(fc_cursor->mcxt);
		fc_cursor->mcxt = NULL;
	}
	fc_arg->ob_type->tp_free(fc_arg);
}

static PyObject * fc_PLy_cursor_iternext(PyObject *fc_self)
{
	PLyCursorObject *fc_cursor;
	PyObject   *fc_ret;
	PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();
	volatile MemoryContext fc_oldcontext;
	volatile ResourceOwner fc_oldowner;
	Portal		fc_portal;

	fc_cursor = (PLyCursorObject *) fc_self;

	if (fc_cursor->closed)
	{
		PLy_exception_set(PyExc_ValueError, "iterating a closed cursor");
		return NULL;
	}

	fc_portal = GetPortalByName(fc_cursor->portalname);
	if (!PortalIsValid(fc_portal))
	{
		PLy_exception_set(PyExc_ValueError,
						  "iterating a cursor in an aborted subtransaction");
		return NULL;
	}

	fc_oldcontext = CurrentMemoryContext;
	fc_oldowner = CurrentResourceOwner;

	PLy_spi_subtransaction_begin(fc_oldcontext, fc_oldowner);

	PG_TRY();
	{
		SPI_cursor_fetch(fc_portal, true, 1);
		if (SPI_processed == 0)
		{
			PyErr_SetNone(PyExc_StopIteration);
			fc_ret = NULL;
		}
		else
		{
			PLy_input_setup_tuple(&fc_cursor->result, SPI_tuptable->tupdesc,
								  fc_exec_ctx->curr_proc);

			fc_ret = PLy_input_from_tuple(&fc_cursor->result, SPI_tuptable->vals[0],
									   SPI_tuptable->tupdesc, true);
		}

		SPI_freetuptable(SPI_tuptable);

		PLy_spi_subtransaction_commit(fc_oldcontext, fc_oldowner);
	}
	PG_CATCH();
	{
		PLy_spi_subtransaction_abort(fc_oldcontext, fc_oldowner);
		return NULL;
	}
	PG_END_TRY();

	return fc_ret;
}

static PyObject * fc_PLy_cursor_fetch(PyObject *fc_self, PyObject *fc_args)
{
	PLyCursorObject *fc_cursor;
	int			fc_count;
	PLyResultObject *fc_ret;
	PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();
	volatile MemoryContext fc_oldcontext;
	volatile ResourceOwner fc_oldowner;
	Portal		fc_portal;

	if (!PyArg_ParseTuple(fc_args, "i:fetch", &fc_count))
		return NULL;

	fc_cursor = (PLyCursorObject *) fc_self;

	if (fc_cursor->closed)
	{
		PLy_exception_set(PyExc_ValueError, "fetch from a closed cursor");
		return NULL;
	}

	fc_portal = GetPortalByName(fc_cursor->portalname);
	if (!PortalIsValid(fc_portal))
	{
		PLy_exception_set(PyExc_ValueError,
						  "iterating a cursor in an aborted subtransaction");
		return NULL;
	}

	fc_ret = (PLyResultObject *) PLy_result_new();
	if (fc_ret == NULL)
		return NULL;

	fc_oldcontext = CurrentMemoryContext;
	fc_oldowner = CurrentResourceOwner;

	PLy_spi_subtransaction_begin(fc_oldcontext, fc_oldowner);

	PG_TRY();
	{
		SPI_cursor_fetch(fc_portal, true, fc_count);

		Py_DECREF(fc_ret->status);
		fc_ret->status = PyLong_FromLong(SPI_OK_FETCH);

		Py_DECREF(fc_ret->nrows);
		fc_ret->nrows = PyLong_FromUnsignedLongLong(SPI_processed);

		if (SPI_processed != 0)
		{
			uint64		fc_i;

			/*
			 * PyList_New()和PyList_SetItem()使用Py_ssize_t作为列表大小
			 * 和列表索引；因此我们无法支持大于
			 * PY_SSIZE_T_MAX的结果。
			 */
			if (SPI_processed > (uint64) PY_SSIZE_T_MAX)
				ereport(ERROR,
						(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
						 errmsg("query result has too many rows to fit in a Python list")));

			Py_DECREF(fc_ret->rows);
			fc_ret->rows = PyList_New(SPI_processed);
			if (!fc_ret->rows)
			{
				Py_DECREF(fc_ret);
				fc_ret = NULL;
			}
			else
			{
				PLy_input_setup_tuple(&fc_cursor->result, SPI_tuptable->tupdesc,
									  fc_exec_ctx->curr_proc);

				for (fc_i = 0; fc_i < SPI_processed; fc_i++)
				{
					PyObject   *fc_row = PLy_input_from_tuple(&fc_cursor->result,
														   SPI_tuptable->vals[fc_i],
														   SPI_tuptable->tupdesc,
														   true);

					PyList_SetItem(fc_ret->rows, fc_i, fc_row);
				}
			}
		}

		SPI_freetuptable(SPI_tuptable);

		PLy_spi_subtransaction_commit(fc_oldcontext, fc_oldowner);
	}
	PG_CATCH();
	{
		PLy_spi_subtransaction_abort(fc_oldcontext, fc_oldowner);
		return NULL;
	}
	PG_END_TRY();

	return (PyObject *) fc_ret;
}

static PyObject * fc_PLy_cursor_close(PyObject *fc_self, PyObject *fc_unused)
{
	PLyCursorObject *fc_cursor = (PLyCursorObject *) fc_self;

	if (!fc_cursor->closed)
	{
		Portal		fc_portal = GetPortalByName(fc_cursor->portalname);

		if (!PortalIsValid(fc_portal))
		{
			PLy_exception_set(PyExc_ValueError,
							  "closing a cursor in an aborted subtransaction");
			return NULL;
		}

		UnpinPortal(fc_portal);
		SPI_cursor_close(fc_portal);
		fc_cursor->closed = true;
	}

	Py_RETURN_NONE;
}
